﻿using System;
using System.Collections.Generic;
using Algoritmia.StateSpaces;
using Algoritmia.DataStructures.PriorityQueues;
using Algoritmia.DataStructures.Sets;

namespace Algoritmia.Schemes.BranchAndBound
{
    public class BestFirstSearcher<TState,TDecision,TWeight,TSolution>
        where TState : class
        where TWeight : IComparable<TWeight>
    {
        private readonly Func<IForwardStateSpace<TState, TDecision>, IPriorityQueue<BranchAndBoundState<TState,TDecision,TWeight>>> _createPriorityQueue;
        private readonly Func<IForwardStateSpace<TState, TDecision>, DataStructures.Sets.ISet<TState>> _createSet;
        private readonly Func<TState, BranchAndBoundDecision<TDecision>, TState, TSolution> _createSolution;

        public BestFirstSearcher(Func<TState, BranchAndBoundDecision<TDecision>, TState, TSolution> createSolution,
                                 Func<IForwardStateSpace<TState, TDecision>, IPriorityQueue<BranchAndBoundState<TState, TDecision, TWeight>>> createPriorityQueue = null,
                                 Func<IForwardStateSpace<TState, TDecision>, DataStructures.Sets.ISet<TState>> createSet = null)
        {
            _createSolution = createSolution;
            _createPriorityQueue = createPriorityQueue ?? (space => new MaxHeap<BranchAndBoundState<TState, TDecision, TWeight>>());
            _createSet = createSet ?? (space => new DummySet<TState>());
        }

        public IEnumerable<TSolution> Enumerate(IForwardStateSpace<TState,TDecision> space, 
            Func<TState,TWeight> priorityFunc = null)
        {
            priorityFunc = priorityFunc ?? (s => default(TWeight));
            var open = _createPriorityQueue(space);
            var visited = _createSet(space);

            foreach(var s in space.InitialStates)
            {
                open.Add(BranchAndBoundState.Create(s, priorityFunc(s), new BranchAndBoundDecision<TDecision>(), s));
                visited.Add(s);
            }

            while (open.Count > 0)
            {
                var s = open.ExtractOptimal();
                var state = s.Current;
                var prev = s.Previous;
                var decision = s.Decision;
                if ( space.IsFinal(state))
                {
                    yield return _createSolution(state, decision, prev);
                }
                foreach(var d in space.Decisions(state))
                {
                    var s1 = space.Decide(state, d);
                    if (!visited.Contains(s1))
                    {
                        open.Add(BranchAndBoundState.Create(s1, priorityFunc(s1), new BranchAndBoundDecision<TDecision>(d, decision), state));
                        visited.Add(s1);
                    }
                }
            }
        }
    }

    public class BestFirstSearcher<TState, TDecision, TWeight> : BestFirstSearcher<TState, TDecision, TWeight, DataStructures.Lists.IList<TDecision>>
        where TState : class
        where TWeight : IComparable<TWeight>
    {
        public BestFirstSearcher(Func<IForwardStateSpace<TState, TDecision>, IPriorityQueue<BranchAndBoundState<TState, TDecision, TWeight>>> createPriorityQueue = null,
                                   Func<IForwardStateSpace<TState, TDecision>, DataStructures.Sets.ISet<TState>> createSet = null)
            : base((s, d, p) => d.RecoverDecisions(), createPriorityQueue, createSet)
        {
        }
    }
}
